
#include "global_avg_pool_tiling.h"
#include "register/op_def_registry.h"

namespace optiling
{
    static ge::graphStatus TilingFunc(gert::TilingContext *context)
    {

        GlobalAvgPoolTilingData tiling;
        uint32_t blockLength;
        uint32_t tileLength;
        uint32_t tileNum;
        uint32_t datatypesize;
        uint32_t outputDim;
        uint32_t lasttileLength;
        auto inputShape = context->GetInputTensor(0)->GetOriginShape();
        uint32_t N = inputShape.GetDim(0);
        uint32_t C = inputShape.GetDim(1);
        outputDim = N * C;//其他维度的值应该是1，你这里有计算其他维度的值吗
        blockLength = context->GetInputTensor(0)->GetShapeSize();
        tileLength = blockLength / outputDim;
        auto dt = context->GetInputTensor(0)->GetDataType();
        if (dt == 0)
        {
            lasttileLength = tileLength % 8;
            tileLength = tileLength / 8 * 8;
        }
        else
        {
            lasttileLength = tileLength % 16;
            tileLength = tileLength / 16 * 16;
        }

        printf("tileLength: %d \n", tileLength);
        printf("lasttileLength: %d \n", lasttileLength);
        printf("outputDim: %d \n", outputDim);
        context->SetBlockDim(1);
        tiling.set_blockLength(blockLength);
        tiling.set_outputDim(outputDim);
        tiling.set_tileLength(tileLength);
        tiling.set_lasttileLength(lasttileLength);
        tiling.set_tileNum(outputDim);
        tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
        context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
        return ge::GRAPH_SUCCESS;
    }
}

namespace ge
{
    static ge::graphStatus InferShape(gert::InferShapeContext *context)
    {
        const gert::Shape *x1_shape = context->GetInputShape(0);
        gert::Shape *y_shape = context->GetOutputShape(0);
        *y_shape = *x1_shape;//这里输出的shape和输入的shape不一致吧
        return GRAPH_SUCCESS;
    }
}

namespace ops
{
    class GlobalAvgPool : public OpDef
    {
    public:
        explicit GlobalAvgPool(const char *name) : OpDef(name)
        {
            this->Input("x")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
            this->Output("y")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

            this->SetInferShape(ge::InferShape);

            this->AICore()
                .SetTiling(optiling::TilingFunc);
            this->AICore().AddConfig("ascend310b");
        }
    };

    OP_ADD(GlobalAvgPool);
}

